from datetime import datetime
from multiprocessing import Process
import random
import threading
import time


_task_registry = []
_num_pool = 4


def sleep_task(*args, **kwargs):
    print(f"start sleep task {args}, {kwargs}")
    time.sleep(random.randint(10, 50))
    print(f"end sleep task {args}, {kwargs}")


def launch_task(task_type, *args, **kwargs):
    if random.random() < 0.3:
        return False, None, "Bad"
    task_id = f"task_{int(time.time() * 1000)}"
    _task_registry.append(
        {
            "task_id": task_id,
            "pid": None,
            "type": task_type,
            "params": kwargs,
            "status": "wait",
            "start": datetime.now(),
            "end": None,
            "result": None,
            "proc": None,
        }
    )
    return True, task_id, "Good"


def cancel_task(task_id, to_status="cancelled"):
    for task in _task_registry:
        if task["task_id"] == task_id and task["status"] in ("running", "wait"):
            if task["status"] == "running":
                task["proc"].terminate()
                task["status"] = to_status
                task["end"] = datetime.now()
                return True
            else:
                task["status"] = to_status
    return False


def rerun_task(task_id):
    for task in _task_registry:
        if task["task_id"] == task_id:
            task["proc"] = None
            task["pid"] = None
            task["result"] = None
            task["status"] = "wait"
            task["start"] = datetime.now()
            task["end"] = None


def get_all_tasks():
    print("get all tasks", _task_registry)
    return [
        {x: y for x, y in task.items() if x not in ("result", "proc")}
        for task in _task_registry
    ]


def get_task_result(task_id):
    print(f"get_task_result", _task_registry)
    for task in _task_registry:
        if task["task_id"] == task_id:
            return task["result"]


def check_task_status():
    while True:
        running_cnt = len([x for x in _task_registry if x["status"] == "running"])
        for task in _task_registry:
            if task["status"] == "running" and task["proc"].is_alive() is False:
                task["status"] = "finished"
                task["end"] = datetime.now()
                task["result"] = {"hahaha": "chenggong"}
            if task["status"] == "wait" and running_cnt < _num_pool:
                proc = Process(target=sleep_task)  # todo
                proc.daemon = True
                proc.start()
                task["proc"] = proc
                task["pid"] = proc.pid
                task["status"] = "running"
        time.sleep(1)


th = threading.Thread(target=check_task_status)
th.daemon = True
th.start()
